import json
import os
import shutil
import time

import numpy as np
import torch
from torch.utils.data import DataLoader
from tqdm import tqdm

from data.neu_seg_competition_nation.dataset import Neu_Seg_Competition_Dataset
from model.UNet import *
from model.model import *
from utils import convert_pred_to_mask, get_scores


def submit_for_ans(unet_model_path,self_model_path):
    # 重新建立 submit 文件夹
    if os.path.exists('./submit'):
        shutil.rmtree('./submit')
    os.mkdir('./submit')
    if not os.path.exists('./submit/baseline_predictions'):
        os.mkdir('./submit/baseline_predictions')
    if not os.path.exists('./submit/test_ground_truths'):
        os.mkdir('./submit/test_ground_truths')
    if not os.path.exists('./submit/test_predictions'):
        os.mkdir('./submit/test_predictions')

    batch_size = 1  # 评估模式，不能动这个
    pin_memory = False
    num_workers = 1
    num_classes = 4
    # 创建设备对象
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")
    unet_model = UNet(3,4)
    unet_model.load_state_dict(torch.load(unet_model_path))
    self_model = self_net()
    self_model.load_state_dict(torch.load(self_model_path))
    # unet_model = torch.load(unet_model_path)
    # self_model = torch.load(self_model_path)

    # 将模型移动到指定设备
    unet_model = unet_model.to(device,non_blocking=True)
    self_model = self_model.to(device,non_blocking=True)

    train_dataset = Neu_Seg_Competition_Dataset(data_path='./data/neu_seg_competition_nation', data_type='training')
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                              pin_memory=pin_memory)
    test_dataset = Neu_Seg_Competition_Dataset(data_path='./data/neu_seg_competition_nation', data_type='test')
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                             pin_memory=pin_memory)
    all_sets = train_dataset + test_dataset
    all_loader = DataLoader(all_sets, batch_size=batch_size, shuffle=False, num_workers=num_workers,
                            pin_memory=pin_memory)
    unet_model.eval()
    self_model.eval()
    shutil.copy2('./model/model.py', './submit/model.py') #复制自建模型结构
    shutil.copy2(self_model_path, './submit/model.pth')    #复制自建模型的参数
    unet_t_all = []
    self_model_t_all = []
    unet_total_u = torch.zeros(num_classes, device=device)
    unet_total_n = torch.zeros(num_classes, device=device)
    self_total_u = torch.zeros(num_classes, device=device)
    self_total_n = torch.zeros(num_classes, device=device)
    with torch.no_grad():
        for data in tqdm(test_loader):
            image, mask, image_names = data
            image_name = image_names[0]
            image = image.to(device)
            mask = mask.to(device)
            self_model_tstart = time.time()
            self_pred = self_model(image)
            self_model_tstop = time.time()
            unet_tstart = time.time()
            unet_pred = unet_model(image)
            unet_tstop = time.time()
            unet_mask = convert_pred_to_mask(unet_pred)
            self_mask = convert_pred_to_mask(self_pred)
            unet_ni, unet_ui = get_scores(unet_pred, mask, device)
            self_ni, self_ui = get_scores(self_pred, mask, device)
            unet_ni = unet_ni.to(device)
            unet_ui = unet_ui.to(device)
            self_ni = self_ni.to(device)
            self_ui = self_ui.to(device)
            unet_total_n += unet_ni
            unet_total_u += unet_ui
            self_total_n += self_ni
            self_total_u += self_ui
            unet_mask_np = unet_mask.cpu().numpy()
            self_mask_np = self_mask.cpu().numpy()
            mask_np = mask.cpu().numpy()
            np.save(f'./submit/baseline_predictions/prediction_{image_name}.npy', unet_mask_np)
            np.save(f'./submit/test_ground_truths/ground_truth_{image_name}.npy', mask_np)
            np.save(f'./submit/test_predictions/prediction_{image_name}.npy', self_mask_np)
            unet_t_all.append(unet_tstop-unet_tstart)
            self_model_t_all.append(self_model_tstop-self_model_tstart)

    unet_total_u = torch.where(unet_total_u < 1, 1, unet_total_u)
    self_total_u = torch.where(self_total_u < 1, 1, self_total_u)
    unet_classiou = unet_total_n / unet_total_u
    self_classiou = self_total_n / self_total_u
    unet_miou = (unet_classiou[1].item() + unet_classiou[2].item() + unet_classiou[3].item()) / 3
    self_miou = (self_classiou[1].item() + self_classiou[2].item() + self_classiou[3].item()) / 3
    unet_fps = 1/(np.mean(unet_t_all))
    self_fps = 1/(np.mean(self_model_t_all))
    unet_parameters = sum(p.numel() for p in unet_model.parameters() if p.requires_grad)
    self_parameters = sum(p.numel() for p in self_model.parameters() if p.requires_grad)
    composite_dict = {
        "UNet": {
            "Class1_IoU": unet_classiou[1].item(),
            "Class2_IoU": unet_classiou[2].item(),
            "Class3_IoU": unet_classiou[3].item(),
            "mIoU": unet_miou,
            "FPS": unet_fps,
            "Parameters": unet_parameters
        },
        "OursModel": {
            "Class1_IoU": self_classiou[1].item(),
            "Class2_IoU": self_classiou[2].item(),
            "Class3_IoU": self_classiou[3].item(),
            "mIoU": self_miou,
            "FPS": self_fps,
            "Parameters": self_parameters
        }
    }
    with open('./submit/关键指标数据文档.txt', 'w') as f:
        json.dump(composite_dict, f)
    formatted_json = json.dumps(composite_dict, indent=4)
    print(formatted_json)

if __name__ == '__main__':
    #测试的是test集
    unet_model_path = './save_model_best/Unet_baseline_miou_0.7565735578536987_model_best_epoch_34 (1).pth'
    self_model_path = './save_model_best/ENet_enhance_b512_random_model_epoch_19999_loss_0.30629223585128784.pth'
    submit_for_ans(unet_model_path,self_model_path)












